package zhengqc.tree;

import java.util.*;

public class HuffmanTree<E extends Comparable<E>> {
    class Node implements Comparable<Node> {
        public Node left;
        public E element;
        public Node right;
        public int weight;
        public Node(Node left, E element, Node right, int weight) {
            this.left = left;
            this.element = element;
            this.right = right;
            this.weight = weight;
        }
        public Node(E element, int weight) {
            this(null, element, null, weight);
        }

        @Override
        public String toString() {
            return String.format("<%d, %d>", this.element, this.weight);
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            Node node = (Node) o;
            return weight == node.weight && Objects.equals(element, node.element);
        }

        @Override
        public int hashCode() {
            return Objects.hash(element, weight);
        }

        @Override
        public int compareTo(Node o) {
            return this.weight - o.weight;
        }

    }
    private Node root;
    private Map<E, String> codeMap;
    private Map<String, E> codeMapReverse;
    private void __init__(Map<E, Integer> map) {
        this.root = null;
        this.codeMap = new TreeMap<>();
        this.codeMapReverse = new TreeMap<>();
        PriorityQueue<Node> pq = new PriorityQueue<>();
        for (E k: map.keySet()) {
            pq.add(new Node(k, map.get(k)));
        }
        Node parent = null;
        while ( pq.size() > 1 ) {
            Node left = pq.poll();
            Node right = pq.poll();
            parent = new Node(left, null, right, left.weight+right.weight);
            pq.add(parent);
        }
        this.root = parent;
        this.backtrace(this.root, "", codeMap);
        for (E k: this.codeMap.keySet()) {
            this.codeMapReverse.put(this.codeMap.get(k), k);
        }
    }
    public HuffmanTree(Map<E, Integer> map) {
        this.__init__(map);
    }
    public HuffmanTree(E[] sequence) {
        Map<E, Integer> map = new TreeMap<>();
        for (E s: sequence) {
            if (map.containsKey(s)) {
                map.put(s, map.get(s)+1);
            } else {
                map.put(s, 1);
            }
        }
        this.__init__(map);
    }
    private void backtrace(Node node, String code, Map<E, String> map) {
        if (node.element != null) {
            map.put(node.element, code);
        }
        if (node.right != null) {
            this.backtrace(node.right, code+'1', map);
        }
        if (node.left != null) {
            this.backtrace(node.left, code+'0', map);
        }
    }
    public Map<E, String> getCodeMap() {
        return this.codeMap;
    }
    public String encode(E[] sequence) {
        StringBuilder sb = new StringBuilder();
        for (E s: sequence) {
            if ( !this.codeMap.containsKey(s) ) {
                throw new IllegalArgumentException(String.format("%s is not exist in code map!", s));
            }
            sb.append(this.codeMap.get(s));
        }
        return sb.toString();
    }
    public ArrayList<E> decode(String code) {
        int i = 0, length = code.length(), beginIndex = 0;
        ArrayList<E> arrayList = new ArrayList<>();
        while ( i < length ) {
            String substr = code.substring(beginIndex, i+1);
            if ( this.codeMapReverse.containsKey( substr ) ) {
                arrayList.add(this.codeMapReverse.get( substr ));
                beginIndex = i + 1;
            }
            i += 1;
        }
        return arrayList;
    }
}
